import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from SwinTransformer_For_Classification import SwinTransformer


# 初始化模型
def initialize_model(num_classes):
    model = SwinTransformer(num_classes=num_classes)
    return model


# 训练模型
def train_model(model, train_loader, criterion, optimizer, num_epochs=10):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * images.size(0)

        epoch_loss = running_loss / len(train_loader.dataset)
        print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}')
    return model


# 保存模型
def save_model(model, path='model.pth'):
    torch.save(model.state_dict(), path)
    print(f'Model saved to {path}')


# 评估模型
def evaluate_model(model, test_loader):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy of the model on the test images: {accuracy}%')
    return accuracy


# 主函数
def main():
    # 数据预处理设置
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # 加载 CIFAR-100 数据集
    train_set = datasets.CIFAR100(root='../Resnet/data', train=True, download=True, transform=transform)
    test_set = datasets.CIFAR100(root='../Resnet/data', train=False, download=True, transform=transform)
    train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

    # 设置模型、损失函数和优化器
    model = initialize_model(num_classes=100)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练和评估模型
    trained_model = train_model(model, train_loader, criterion, optimizer, num_epochs=10)
    save_model(trained_model, '../pt_file/SwinTransformer_classification_train.pt')
    evaluate_model(trained_model, test_loader)


if __name__ == "__main__":
    main()
